#/usr/bin/python27
#-*- coding: utf-8 -*-
#  Copyright 2011-2013 wuweihui@cn.ibm.com
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import telnetlib
import time
import re
import select

def _close_log(msg, level=None):
   pass
   
   
class TelnetConnection(telnetlib.Telnet):

   def __init__(self, host=None, port=23, timeout=3.0, newline='CRLF',
                 prompt=None, prompt_is_regexp=False):
      port = port == '' and 23 or int(port)
      telnetlib.Telnet.__init__(self, host, port)
      self.set_timeout(timeout)
      self.set_newline(newline)
      self.set_prompt(prompt, prompt_is_regexp)
      self._default_log_level = 'INFO'
      self.set_option_negotiation_callback(self._negotiate_echo_on)
      
   def set_timeout(self, timeout):
      """Sets the timeout used in read operations to the given value.
      
        `timeout` is given in Robot Framework's time format
        (e.g. 1 minute 20 seconds) that is explained in the User Guide.
      
        Read operations that expect some output to appear (`Read
        Until`, `Read Until Regexp`, `Read Until Prompt`) use this
        timeout and fail if the expected output has not appeared when
        this timeout expires.
      
        The old timeout is returned and can be used to restore it later.
      
        Example:
        | ${tout} = | Set Timeout | 2 minute 30 seconds |
        | Do Something |
        | Set Timeout | ${tout} |
        """
      old = getattr(self, '_timeout', 3.0)
      #print timeout
      self._timeout = float(timeout)
      return old
      
   def set_newline(self, newline):
      """Sets the newline used by the `Write` keyword.
      
        Newline can be given either in escaped format using '\\n' and
        '\\r', or with special 'LF' and 'CR' syntax.
      
        Examples:
        | Set Newline | \\n  |
        | Set Newline | CRLF |
      
        Correct newline to use depends on the system and the default
        value is 'CRLF'.
      
        The old newline is returned and can be used to restore it later.
        See `Set Prompt` or `Set Timeout` for an example.
        """
      old = getattr(self, '_newline', 'CRLF')
      self._newline = newline.upper().replace('LF','\n').replace('CR','\r')
      return old
      
   def close_connection(self, loglevel=None):
      """Closes the current Telnet connection and returns any remaining output.
      
        See `Read` for more information on `loglevel`.
        """
      telnetlib.Telnet.close(self)
      ret = self.read_all().decode('ASCII', 'ignore')
      self._log(ret, loglevel)
      return ret
      
   def login(self, username, password, login_prompt='login: ',
              password_prompt='Password: '):
      """Logs in to the Telnet server with the given user information.
      
        The login keyword reads from the connection until login_prompt
        is encountered and then types the user name. Then it reads
        until password_prompt is encountered and types the
        password. The rest of the output (if any) is also read, and
        all the text that has been read is returned as a single
        string.
      
        If a prompt has been set to this connection, either with `Open
        Connection` or `Set Prompt`, this keyword reads the output
        until the prompt is found. Otherwise, the keyword sleeps for a
        second and reads everything that is available.
        """
      ret = self.read_until(login_prompt, 'TRACE').decode('ASCII', 'ignore')
      self.write_bare(username + self._newline)
      ret += username + '\n'
      ret += self.read_until(password_prompt, 'TRACE').decode('ASCII', 'ignore')
      self.write_bare(password + self._newline)
      ret += '*' * len(password) + '\n'
      if self._prompt_is_set():
         try:
            ret += self.read_until_prompt('TRACE')
         except AssertionError:
            self._verify_login(ret)
            raise
      else:
         ret += self._verify_login(ret)
      self._log(ret)
      return ret
      
   def _verify_login(self, ret):
   # It is necessary to wait for the 'login incorrect' message to appear.
      time.sleep(1)
      while True:
         try:
            ret += self.read_until('\n', 'TRACE').decode('ASCII', 'ignore')
         except AssertionError:
            return ret
         else:
            if 'Login incorrect' in ret:
               self._log(ret)
               raise AssertionError("Login incorrect")
               
   def write(self, text, loglevel=None):
      """Writes the given text over the connection and appends a newline.
      
        Consumes the written text (until the appended newline) from
        the output and returns it. The given text must not contain
        newlines.
      
        Note: This keyword does not return the possible output of the
        executed command. To get the output, one of the `Read XXX`
        keywords must be used.
      
        See `Read` for more information on `loglevel`.
        """
      if self._newline in text:
         raise RuntimeError("Write cannot be used with string containing "
            "newlines. Use 'Write Bare' instead.")
      text += self._newline
      self.write_bare(text)
      # Can't read until 'text' because long lines are cut strangely in the output
      return self.read_until(self._newline, loglevel)
      
   def write_bare(self, text):
      """Writes the given text over the connection without appending a newline.
      
        Does not consume the written text.
        """
      try:
         text = str(text)
      except UnicodeError:
         raise ValueError('Only ASCII characters are allowed in Telnet. '
            'Got: %s' % text)
      telnetlib.Telnet.write(self, text)
      
   def write_until_expected_output(self, text, expected, timeout,
                                    retry_interval, loglevel=None):
      """Writes the given text repeatedly, until `expected` appears in the output.
      
        `text` is written without appending a
        newline. `retry_interval` defines the time waited before
        writing `text` again. `text` is consumed from the output
        before `expected` is tried to be read.
      
        If `expected` does not appear in the output within `timeout`,
        this keyword fails.
      
        See `Read` for more information on `loglevel`.
      
        Example:
        | Write Until Expected Output | ps -ef| grep myprocess\\n | myprocess |
        | ...                         | 5s                        | 0.5s      |
      
        This writes the 'ps -ef | grep myprocess\\n', until
        'myprocess' appears on the output. The command is written
        every 0.5 seconds and the keyword ,fails if 'myprocess' does
        not appear in the output in 5 seconds.
        """
      timeout = int(timeout)
      retry_interval = float(retry_interval)
      starttime = time.time()
      while time.time() - starttime < timeout:
         self.write_bare(text)
         self.read_until(text, loglevel)
         ret = telnetlib.Telnet.read_until(self, expected,
                                              retry_interval).decode('ASCII', 'ignore')
         self._log(ret, loglevel)
         if ret.endswith(expected):
            return ret
      raise AssertionError("No match found for '%s' in %s"
         % (expected, str(timeout)))
         
   def read(self, loglevel=None):
      """Reads and returns/logs everything that is currently available in the output.
      
        The read message is always returned and logged. The default
        log level is either 'INFO', or the level set with `Set Default
        Log Level`.  `loglevel` can be used to override the default
        log level, and the available levels are TRACE, DEBUG, INFO,
        and WARN.
        """
      ret = self.read_very_eager().decode('ASCII', 'ignore')
      self._log(ret, loglevel)
      return ret
      
   def read_until(self, expected, loglevel=None):
      """Reads from the current output, until expected is encountered.
      
        Text up to and including the match is returned. If no match is
        found, the keyword fails.
      
        See `Read` for more information on `loglevel`.
        """
      ret = telnetlib.Telnet.read_until(self, expected,
                                          self._timeout).decode('ASCII', 'ignore')
      self._log(ret, loglevel)
      if not ret.endswith(expected):
         #print repr(ret)
         raise AssertionError("No match found for '%s' in %s"
            % (expected, str(self._timeout)))
      return ret
      
   def read_until_regexp(self, *expected):
      """Reads from the current output, until a match to a regexp in expected.
      
        Expected is a list of regular expression patterns as strings,
        or compiled regular expressions. The keyword returns the text
        up to and including the first match to any of the regular
        expressions.
      
        If the last argument in `*expected` is a valid log level, it
        is used as `loglevel` in the keyword `Read`.
      
        Examples:
        | Read Until Regexp | (#|$) |
        | Read Until Regexp | first_regexp | second_regexp |
        | Read Until Regexp | some regexp  | DEBUG |
        """
      expected = list(expected)
      if self._is_valid_log_level(expected[-1]):
         loglevel = expected[-1]
         expected = expected[:-1]
      else:
         loglevel = 'INFO'
      try:
         index, _, ret = self.expect(expected, self._timeout)
      except TypeError:
         index, ret = -1, ''
      ret = ret.decode('ASCII', 'ignore')
      self._log(ret, loglevel)
      if index == -1:
         expected = [ exp if isinstance(exp, basestring) else exp.pattern
            for exp in expected ]
         raise AssertionError("No match found for %s in %s"
            % (str(expected),
                                    str(self._timeout)))
      return ret
      
   def read_until_prompt(self, loglevel=None):
      """Reads from the current output, until a prompt is found.
      
        The prompt must have been set, either in the library import or
        at login time, or by using the `Set Prompt` keyword.
      
        See `Read` for more information on `loglevel`.
        """
      if not self._prompt_is_set():
         raise RuntimeError('Prompt is not set')
      prompt, regexp = self._prompt
      if regexp:
         return self.read_until_regexp(prompt, loglevel)
      return self.read_until(prompt, loglevel)
      
   def execute_command(self, command, loglevel=None):
      """Executes given command and reads and returns everything until prompt.
      
        This is a convenience keyword; following two are functionally
        identical:
      
        | ${out} = | Execute Command   | Some command |
      
        | Write    | Some command      |
        | ${out} = | Read Until Prompt |
      
        This keyword expects a prompt to be set, see `Read Until
        Prompt` for details.
      
        See `Read` for more information on `loglevel`.
        """
      self.write(command, loglevel)
      return self.read_until_prompt(loglevel)
      
   def set_prompt(self, prompt, prompt_is_regexp=False):
      """Sets the prompt used in this connection to `prompt`.
      
        If `prompt_is_regexp` is a non-empty string, the given prompt is
        considered to be a regular expression.
      
        The old prompt is returned and can be used to restore it later.
      
        Example:
        | ${prompt} | ${regexp} = | Set Prompt | $ |
        | Do Something |
        | Set Prompt | ${prompt} | ${regexp} |
        """
      old = hasattr(self, '_prompt') and self._prompt or (None, False)
      if prompt_is_regexp:
         self._prompt = (re.compile(prompt), True)
      else:
         self._prompt = (prompt, False)
      if old[1]:
         return old[0].pattern, True
      return old
      
   def _prompt_is_set(self):
      return self._prompt[0] is not None
      
   def expect(self, list, timeout=None):
        re = None
        indices = range(len(list))
        if timeout is not None:
            from time import time
            time_start = time()
        while 1:
            self.process_rawq()
            pos = max(0, len(self.cookedq) - 100)
            for i in indices:
                m = list[i].search(self.cookedq, pos)
                if m:
                    text = self.cookedq[:m.end()]
                    self.cookedq = self.cookedq[m.end():]
#                    print "*DEBUG* |%s| %d %d" % (m.string[m.start():m.end()], m.end(), m.endpos)
                    return (i, m, text)
            if self.eof:
                self._log("Telnet: Eof detected", "TRACE")
                break
            if timeout is not None:
                elapsed = time() - time_start
                if elapsed >= timeout:
                    self._log("Telnet: Elapsed time exceeds timeout -> No further check", "TRACE")
                    break
                s_args = ([self.fileno()], [], [], timeout-elapsed)
                r, w, x = select.select(*s_args)
                if not r:
                    self._log("Telnet: Select System call timed out", "TRACE")
                    break
            self.fill_rawq()
        text = self.read_very_lazy()
        if not text and self.eof:
            raise EOFError
        return (-1, None, text)

   def process_rawq(self):
        """Transfer from raw queue to cooked queue.

        Set self.eof when connection is closed.  Don't block unless in
        the midst of an IAC sequence.

        """
        buf = ['', '']
        try:
            while self.rawq:
                c = self.rawq_getchar()
                if not self.iacseq:
                    if c == telnetlib.theNULL:
                        continue
                    if c == "\021":
                        continue
                    if c != telnetlib.IAC:
                        buf[self.sb] = buf[self.sb] + c
                        continue
                    else:
                        self.iacseq += c
                elif len(self.iacseq) == 1:
                    # 'IAC: IAC CMD [OPTION only for WILL/WONT/DO/DONT]'
                    if c in (telnetlib.DO, telnetlib.DONT, telnetlib.WILL, telnetlib.WONT):
                        self.iacseq += c
                        continue

                    self.iacseq = ''
                    if c == telnetlib.IAC:
                        buf[self.sb] = buf[self.sb] + c
                    else:
                        if c == telnetlib.SB: # SB ... SE start.
                            self.sb = 1
                            self.sbdataq = ''
                        elif c == telnetlib.SE:
                            self.sb = 0
                            self.sbdataq = self.sbdataq + buf[1]
                            buf[1] = ''
                        if self.option_callback:
                            # Callback is supposed to look into
                            # the sbdataq
                            self.option_callback(self.sock, c, telnetlib.NOOPT)
                        else:
                            # We can't offer automatic processing of
                            # suboptions. Alas, we should not get any
                            # unless we did a WILL/DO before.
                            self.msg('IAC %d not recognized' % ord(c))
                elif len(self.iacseq) == 2:
                    cmd = self.iacseq[1]
                    self.iacseq = ''
                    opt = c
                    if cmd in (telnetlib.DO, telnetlib.DONT):
                        self.msg('IAC %s %d',
                            cmd == telnetlib.DO and 'DO' or 'DONT', ord(opt))
                        if self.option_callback:
                            self.option_callback(self.sock, cmd, opt)
                        else:
                            self.sock.sendall(telnetlib.IAC + telnetlib.WONT + opt)
                    elif cmd in (telnetlib.WILL, telnetlib.WONT):
                        self.msg('IAC %s %d',
                            cmd == telnetlib.WILL and 'WILL' or 'WONT', ord(opt))
                        if self.option_callback:
                            self.option_callback(self.sock, cmd, opt)
                        else:
                            self.sock.sendall(telnetlib.IAC + telnetlib.DONT + opt)
        except EOFError: # raised by self.rawq_getchar()
            self.iacseq = '' # Reset on EOF
            self.sb = 0
            pass
        i = 0;
        while i < len(buf[0]):
            if buf[0][i] == telnetlib.NAOL:        #process BS
                if i == 0:                       
                    buf[0] = buf[0][1:]             
                    if len(self.cookedq) != 0 and (self.cookedq[-1] != chr(10) and self.cookedq[-1] != chr(13)):
                         self.cookedq = self.cookedq[:-1]
                else:
                    if buf[0][i-1] == chr(10) or buf[0][i-1] == chr(13):
                        buf[0] = buf[0][:i] + buf[0][i + 1:]       
                    else:
                        buf[0] = buf[0][:i-1] + buf[0][i + 1:]     
                        i = i - 1
            elif buf[0][i] == telnetlib.OUTMRK:    #process OUTMARK
                #print 'outmark ++++', buf[0][i:]
                match = re.match(r'\\[\d+\;\d+H', buf[0][i:])          #new line
                if match:
                    outmark = match.group()
                    buf[0] = buf[0][:i] + self._newline + buf[0][i + len(outmark):]
                    i = i + len(self._newline)
                else:
                    i = i + 1
            else:
                i = i + 1
        self.cookedq = self.cookedq + buf[0]
        self.sbdataq = self.sbdataq + buf[1]

      
   def set_default_log_level(self, level):
      """Sets the default log level used by all read keywords.
      
        The possible values are TRACE, DEBUG, INFO and WARN. The default is
        INFO. The old value is returned and can be used to restore it later,
        similarly as with `Set Timeout`.
        """
      self._is_valid_log_level(level, raise_if_invalid=True)
      old = self._default_log_level
      self._default_log_level = level.upper()
      return old
      
   def _log(self, msg, level=None):
      self._is_valid_log_level(level, raise_if_invalid=True)
      msg = msg.strip()
      if level is None:
         level = self._default_log_level
      if msg != '':
         print '*%s* %s' % (level.upper(), msg)
         
   def _is_valid_log_level(self, level, raise_if_invalid=False):
      if level is None:
         return True
      if isinstance(level, basestring) and \
                level.upper() in ['TRACE', 'DEBUG', 'INFO', 'WARN']:
         return True
      if not raise_if_invalid:
         return False
      raise AssertionError("Invalid log level '%s'" % level)
      
   def _negotiate_echo_on(self, sock, cmd, opt):
   # This is supposed to turn server side echoing on and turn other options off.
      if opt == telnetlib.ECHO and cmd in (telnetlib.WILL, telnetlib.WONT):
         self.sock.sendall(telnetlib.IAC + telnetlib.DO + opt)
      elif opt != telnetlib.NOOPT:
         if cmd in (telnetlib.DO, telnetlib.DONT):
            self.sock.sendall(telnetlib.IAC + telnetlib.WONT + opt)
         elif cmd in (telnetlib.WILL, telnetlib.WONT):
            self.sock.sendall(telnetlib.IAC + telnetlib.DONT + opt)
   
   def close_log(self):
      old = self._log
      self._log = _close_log
      return old
  
            
class StorageTelnet(TelnetConnection):
   def __init__(self, host=None, port=23, timeout=5.0, newline='CRLF', computername=None):
      port = port == '' and 23 or int(port)
      TelnetConnection.__init__(self, host, port, timeout, newline)
      self.computername = computername
      
   def login(self, username, password, login_prompt='login: ',
              password_prompt='Password: '):
      if self.computername:
         prompt = '\[%s@%s\s[~a-zA-Z_0-9\-]+\]# ' % (username, self.computername)
      else:
         prompt = '\[%s@[a-zA-Z_0-9\-]+\s[~a-zA-Z_0-9\-]+\]# ' % (username)
      self.set_prompt(prompt, True)
      
      #super(uTelnet, self).login(username, password, login_prompt, password_prompt)
      TelnetConnection.login(self, username, password, login_prompt, password_prompt)
      
class WinTelnet(TelnetConnection):
   def __init__(self, host=None, port=23, timeout=15.0, newline='CRLF'):
      port = port == '' and 23 or int(port)
      TelnetConnection.__init__(self, host, port, timeout, newline)
      self.set_option_negotiation_callback(self._negotiate_terminal)
      
   def _negotiate_terminal(self, sock, cmd, opt):
       if opt == telnetlib.ECHO and cmd in (telnetlib.WILL, telnetlib.WONT):
         self.sock.sendall(telnetlib.IAC + telnetlib.DO + opt)
       elif cmd == telnetlib.DO and opt == telnetlib.TTYPE:
           sock.sendall(telnetlib.IAC + telnetlib.WILL + opt)
       elif cmd == telnetlib.DO and opt == telnetlib.NAWS:       #negotiate winsize
           sock.sendall(telnetlib.IAC + telnetlib.WILL + opt)
           nopt = chr(0x1)+chr(0xF4)+chr(0x4E)+chr(0x20)+chr(0xff)+chr(0xf0)   #500*20000
           sock.sendall(telnetlib.IAC + telnetlib.SB + nopt)
       elif cmd == telnetlib.WILL and opt == telnetlib.SGA:
           self.sock.sendall(telnetlib.IAC + telnetlib.DO + opt)
       elif opt != telnetlib.NOOPT:
         if cmd in (telnetlib.DO, telnetlib.DONT):
            self.sock.sendall(telnetlib.IAC + telnetlib.WONT + opt)
         elif cmd in (telnetlib.WILL, telnetlib.WONT):
            self.sock.sendall(telnetlib.IAC + telnetlib.DONT + opt)
       elif cmd in [telnetlib.SE] and self.sbdataq == chr(0x18)+chr(0x01):
           opt = chr(0x18)+chr(0x00)+chr(0x41)+chr(0x4e)+chr(0x53)+chr(0x49)+chr(0xff)+chr(0xf0)
           self.sock.sendall(telnetlib.IAC + telnetlib.SB + opt)
       
    
   def login(self, username, password, login_prompt='login: ',
              password_prompt='password: '):
      prompt = '[A-Za-z]\:\\\\[\w \\\\\-\.]*>'
      self.set_prompt(prompt, True)
      return TelnetConnection.login(self, username, password, login_prompt, password_prompt)

class ScriptLine(object):
   
   re_level = re.compile(r'\s*(\d+)\-')
   
   def __init__(self, srcline = ''):
      self.restorelevel = 0
      self.depth = 0
      self.script = ''
      if srcline:
         self.process_src_line(srcline)
   
   def process_src_line(self, srcline):
      ret = self.re_level.match(srcline)
      if ret:
         self.restorelevel = ret.group()
      for s in srcline:
         if s == ' ':
            self.depth += 1
      self.script = srcline.strip()
      if ret:
         self.script = self.script[len(self.restorelevel):]

class Script(object):
   
   re_precondition = re.compile(r'\s*\[precondition\]')
   re_conditionrestore= re.compile(r'\s*\[conditionrestore\]')
   re_case = re.compile(r'\s*\[case\]')
   re_caserestore = re.compile(r'\s*\[caserestore\]')
   
   PRECONDITION = 0
   PRECONDITIONRESTOR = 1
   CASE = 2
   CASERESTORE = 3
   
   def __init__(self, fn=None):
      if fn:
         self.read(fn)
   
   def read(self, fn):
      self.precondition = []    #[[restore level, depth, script line]...]
      self.precondtionrestore = []
      self.case = []
      self.caserestore = []
      
      self.fn = fn
      with open(fn, 'rb') as fp:
         src = fp.read()
      scripttype = 0
      for line in src.splitlines():
         if not line.strip():
            continue
         
         if self.re_precondition.match(line.lower()):
            scripttype = self.PRECONDITION
            continue
         if self.re_conditionrestore.match(line.lower()):
            scripttype = self.PRECONDITIONRESTOR
            continue
         if self.re_case.match(line.lower()):
            scripttype = self.CASE
            continue
         if self.re_caserestore.match(line.lower()):
            scripttype = self.CASERESTORE
            continue         
         
         scriptline = ScriptLine(line)
               
         if scripttype == self.PRECONDITION:
            self.precondition.append(scriptline)
         elif scripttype == self.PRECONDITIONRESTOR:
            self.precondtionrestore.append(scriptline)
         elif scripttype == self.CASE:
            self.case.append(scriptline)
         elif scripttype == self.CASERESTORE:
            self.caserestore.append(scriptline)
         
   
   def run(self, func, errorhandle=None):
      #setup
      try:
         for scriptline in self.precondition:
            func(scriptline.script)
      except Exception, e:
         try:
            for restoreline in self.precondtionrestore:
               if restoreline.restorelevel < scriptline.restorelevel:
                  func(restoreline.script)
         except Exception, e:
            if errorhandle: errorhandle(e)
         
         return 1
      
      #run case
      casecount = len(self.case)
      i = 0
      caserestoreflags = []
      while i < casecount:
         try:
            func(self.case[i].script)
            pass
         except Exception, e:
            if errorhandle: errorhandle(e)
            if self.case[i].restorelevel != 0: caserestoreflags.append(self.case[i].restorelevel)
            for j in range(i+1, casecount):
               if self.case[j].depth <= self.case[i].depth:
                  i = j
                  continue
               if self.case[j].restorelevel != 0: caserestoreflags.append(self.case[j].restorelevel)
         i += 1
      
#      try:
#         for scriptline in self.case:
#            func(scriptline.script)
#      except Exception, e:
#         if errorhandle: errorhandle(e)
#         #restore case
#         try:
#            for restoreline in self.caserestore:
#               if restoreline.restorelevel < scriptline.restorelevel:
#                  func(restoreline.script)
#         except Exception, e:
#            if errorhandle: errorhandle(e)
#         try:
#            #restore setup
#            for restoreline in self.precondtionrestore:
#               func(restoreline.script)
#         except Exception, e:
#            if errorhandle: errorhandle(e)
#         return 2
      
      #normal restore
      try:
         for restoreline in self.caserestore:
            if restoreline.restorelevel not in caserestoreflags:
               func(restoreline.script)
      except Exception, e:
         if errorhandle: errorhandle(e)
      #restore setup
      try:
         for restoreline in self.precondtionrestore:
            func(restoreline.script)
      except Exception, e:
         if errorhandle: errorhandle(e)
         return 3
      
      return 0
   
   def run_files(self, *files):
      pass
   


def test_print(arg):
   print str(arg)

if __name__ == '__main__':
   t = Script('script.sc')
   t.run(test_print)
   
      